import numpy as np

class BallTreeNode:
    """定义BallTree的节点结构"""
    def __init__(self, points, indices):
        self.points = points  # 当前节点包含的数据点
        self.indices = indices  # 数据点的索引
        self.center = np.mean(points, axis=0)  # 节点中心
        self.radius = np.max(np.linalg.norm(points - self.center, axis=1))  # 包围球半径
        self.left = None  # 左子节点
        self.right = None  # 右子节点

class BallTree:
    """BallTree实现"""
    def __init__(self, data, leaf_size=2):
        self.data = data
        self.leaf_size = leaf_size  # 叶子节点的最小大小
        self.root = self._build_tree(data, np.arange(len(data)))
    
    def _build_tree(self, points, indices):
        """
        递归构建BallTree
        :param points: 当前节点包含的点
        :param indices: 点的索引
        :return: 构建的节点
        """
        # 如果点数量小于叶子节点大小，则创建叶子节点
        if len(points) <= self.leaf_size:
            return BallTreeNode(points, indices)
        
        # 找到两个最远点，进行划分
        p1, p2 = self._find_farthest_points(points)
        left_indices = []
        right_indices = []
        
        for idx, point in enumerate(points):
            if np.linalg.norm(point - points[p1]) < np.linalg.norm(point - points[p2]):
                left_indices.append(idx)
            else:
                right_indices.append(idx)
        
        left_points = points[left_indices]
        right_points = points[right_indices]
        
        # 创建当前节点
        node = BallTreeNode(points, indices)
        node.left = self._build_tree(left_points, indices[left_indices])
        node.right = self._build_tree(right_points, indices[right_indices])
        return node
    
    def _find_farthest_points(self, points):
        """
        找到点集中距离最远的两个点
        :param points: 点集
        :return: 最远的两个点的索引
        """
        max_distance = -1
        p1, p2 = 0, 0
        
        for i in range(len(points)):
            for j in range(i + 1, len(points)):
                distance = np.linalg.norm(points[i] - points[j])
                if distance > max_distance:
                    max_distance = distance
                    p1, p2 = i, j
        return p1, p2

# 数据生成
np.random.seed(42)
data = np.random.rand(10, 3)  # 生成10个三维数据点

# 构建BallTree
tree = BallTree(data)

# 打印节点信息
def print_tree(node, depth=0):
    if node is None:
        return
    print(f"{'  ' * depth}节点中心: {node.center}, 半径: {node.radius}, 包含点数: {len(node.points)}")
    print_tree(node.left, depth + 1)
    print_tree(node.right, depth + 1)

print("BallTree结构:")
print_tree(tree.root)